import torch
import cv2 as cv
import numpy as np
import torch.nn.functional as F
from unet import UNet
from train import *


def imshow(name, label):
    m, n = label.shape
    temp = np.zeros((m, n, 3))
    temp[label == 1] = [255, 0, 0]
    temp[label == 2] = [0, 255, 0]
    temp[label == 3] = [0, 0, 255]
    cv.imshow(name, temp)
    cv.waitKey(1)


if __name__ == '__main__':
    model_path = find_last_pt("checkpoints/")
    image_path = "../datasets/images/wd02.jpg"
    size = (512, 512)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    model = UNet(n_channels=3, n_classes=4).to(device)
    checkpoint = torch.load(model_path, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()

    img = cv.imread(image_path, -1)
    img = cv.resize(img, size, interpolation=cv.INTER_NEAREST)
    img = img[:, :, ::-1]
    img = img.transpose((2, 0, 1))
    img = img.astype(np.float32) / 255.0
    img = torch.from_numpy(img)
    img = img.unsqueeze(0)
    img = img.to(device)

    with torch.no_grad():
        output = model(img)
        probs = F.softmax(output, dim=1)
        probs = probs.squeeze(0).cpu().numpy()
    dst = np.argmax(probs, axis=0).astype(np.uint8)

    imshow("det", dst)
    cv.waitKey(0)
    cv.destroyAllWindows()
